import logging
import os

import streamlit as st
from llama_index.core import Settings
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)

# 1. 配置阿里云API（替换成您的API Key）
print('初始化模型对象')
llm_dash = DashScope(
    model="qwen-max",  # 使用通义千问Max模型 qwen-max qwen-turbo
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3,
    system_prompt="""作为一个虚拟对话助手，我将尝试与用户聊天。
          同时，在用户需要时，帮助用户提出问题并回答关于下一步的问题。
          仅根据提供的文档内容进行回答,如果话题与文档内容不相关,回答:抱歉这个话题我不擅长。
          """
)

# 2. 使用本地嵌入模型（节省成本）
print('使用本地嵌入模型（节省成本）')
embed_model_zh = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)


def load_data():
    data_dir = 'D:/pythonRoom/llm_llama/data2'
    reader = SimpleDirectoryReader(input_dir=data_dir, recursive=True)
    docs = reader.load_data()
    # 配置全局设置
    Settings.embed_model = embed_model_zh
    Settings.llm = llm_dash
    index = VectorStoreIndex.from_documents(docs)
    return index


vector_index = load_data()

'''
streamlit run class03/p3_3_2streamlit.py
'''
st.markdown("这是一段Markdown文本")
st.title('这是一个标题')
st.header('这是一个一级标题')
st.subheader('这是一个二级标题')
st.caption('这是一段解释上面内容的说明文字')

prompt = st.chat_input("说点什么")
if prompt:
    st.write(f"用户发送了以下提示：{prompt}")

st.session_state.messages = [
    {"role": "user", "content": "OPPO手机中最贵的是哪一款？"}
]

query_engine = vector_index.as_query_engine()
# 初始化一个ChatEngine实例并保存到session_state中
if "chat_engine" not in st.session_state.keys():
    chat_engine = CondenseQuestionChatEngine.from_defaults(
        query_engine=query_engine
    )
    st.session_state.chat_engine = chat_engine
